import abc
from dataclasses import dataclass

from zkl_aiutils_training import Resumable, ResumableController
from zkl_pyutils_fsspec import Fs, resolve_child_fs


class FsMixin(abc.ABC):
    @property
    @abc.abstractmethod
    def fs(self) -> Fs | None:
        pass

    @property
    def checkpoints_fs(self) -> Fs | None:
        return resolve_child_fs(self.fs, "checkpoints") \
            if self.fs is not None else None

    def get_checkpoint_fs(self, checkpoint_key: str) -> Fs:
        return resolve_child_fs(self.checkpoints_fs, checkpoint_key) \
            if self.checkpoints_fs is not None else None


class FsCheckpointMixin(FsMixin, abc.ABC):
    @property
    @abc.abstractmethod
    def checkpoint_key(self) -> str | None:
        pass

    @property
    def checkpoint_fs(self) -> Fs | None:
        return self.get_checkpoint_fs(self.checkpoint_key) \
            if self.checkpoint_key is not None else None


@dataclass(kw_only=True)
class FsResumeArgs(FsMixin, abc.ABC):
    fs: Fs | None = None


@dataclass(kw_only=True)
class FsResumeFromScratchArgs(FsResumeArgs):
    fs: Fs | None = None


@dataclass(kw_only=True)
class FsResumeFromCheckpointArgs(FsResumeArgs, FsCheckpointMixin):
    fs: Fs | None = None
    checkpoint_key: str | None = None


@dataclass(kw_only=True)
class FsPauseArgs(FsCheckpointMixin):
    fs: Fs | None = None
    checkpoint_key: str | None = None


class FsResumable(Resumable[FsResumeArgs, FsPauseArgs], abc.ABC):
    def on_resume(self, args: FsResumeArgs):
        pass

    def on_pause(self, args: FsPauseArgs):
        pass


class CheckpointKeyResolver(abc.ABC):
    @abc.abstractmethod
    def get_latest_checkpoint_key(self, args: FsResumeArgs) -> str:
        pass

    @abc.abstractmethod
    def get_current_checkpoint_key(self) -> str:
        pass


class FsResumableController(ResumableController[FsResumeArgs, FsPauseArgs], FsMixin):
    _fs: Fs | None = None

    @property
    def fs(self) -> Fs | None:
        return self._fs

    # resume

    def resume(self, args: FsResumeArgs):
        args = self._resolve_resume_args(args)
        self._resume(args)

    def _resolve_resume_args(self, args: FsResumeArgs) -> FsResumeArgs:
        if not isinstance(args, FsResumeArgs):
            raise TypeError(f"Expected FsResumeArgs, got {args}")
        if isinstance(args, FsResumeFromCheckpointArgs) and args.checkpoint_key is None:
            resolver = self.dock.get_plugin(CheckpointKeyResolver)
            args.checkpoint_key = resolver.get_latest_checkpoint_key(args)
        return args

    def _resume(self, args: FsResumeArgs):
        args.fs.makedirs("", exist_ok=True)
        super()._resume(args)
        self._fs = args.fs

    # pause

    def pause(self, args: FsPauseArgs | None = None):
        args = self._resolve_pause_args(args)
        self._pause(args)

    def _resolve_pause_args(self, args: FsPauseArgs | None) -> FsPauseArgs:
        if args is None:
            args = FsPauseArgs()
        if not isinstance(args, FsPauseArgs):
            raise TypeError(f"Expected FsPauseArgs, got {args}")
        if args.fs is None:
            if self._fs is None:
                raise RuntimeError("Cannot resolve fs for pausing!")
            args.fs = self._fs
        if args.checkpoint_key is None:
            resolver = self.dock.get_plugin(CheckpointKeyResolver)
            args.checkpoint_key = resolver.get_current_checkpoint_key()
        return args

    def _pause(self, args: FsPauseArgs):
        args.checkpoint_fs.makedirs("", exist_ok=True)
        super()._pause(args)
